# Copyright 2021-2024 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Rigid Body System
"""

import sys
import math
import numpy as np
sys.path.append('../../../')

import torch
from torch import Tensor, nn
import torch.nn.functional as F
from torch.nn import Parameter

from ..molecule import Molecule
from ...function.units import GLOBAL_UNITS, Units
from ..modelling.mol2_parser import mol2parser

PI = 3.1415927
ATOM_MASS = {'C': 12.01, 'H': 1.008}


class BenzRigidBody(Molecule):
    """ Class for Benzene toy model system.
    Args:
        mol2(str): The input mol2 format file name.
        length_unit(str): The input length unit.
        residues: TODO
        index(ndarray): The rigid atom indexes in the system.
    """
    def __init__(self, mol2, length_unit=None, residues=None, index=None):
        super().__init__()
        if length_unit is None:
            length_unit = GLOBAL_UNITS.length_unit
        self.units = Units(length_unit, GLOBAL_UNITS.energy_unit)
        self.einsum_0 = lambda x, y: torch.einsum('ij,ij->i', x, y)
        self.concat = lambda *args, axis=-2: torch.cat(args, dim=axis)
        self.L1 = 1.4*math.sqrt(3)
        self.L2 = 1.4*math.sqrt(3)
        self.Angle = PI / 3
        self.Hight = self.L2 * math.sin(self.Angle)
        self.Floor = math.sqrt(self.L2 ** 2 - self.Hight ** 2)
        mol2_obj = mol2parser(mol2)
        # (A, )
        atom_names = mol2_obj.atom_names
        self.atom_name = [atom_names]
        # (B, A)
        self.atom_type = np.array([[atype.replace('.AR', 'A').replace('H', 'HA') for atype in mol2_obj.atom_types]])
        charges = mol2_obj.charges
        self.atom_charge = torch.tensor(charges, dtype=torch.float32)
        res_ids = mol2_obj.res_ids
        self.atom_resids = torch.tensor(res_ids, dtype=torch.int32)
        self.atom_mass = torch.tensor([ATOM_MASS[name] for name in atom_names], dtype=torch.float32)
        self.atom_mask = torch.ones(self.atom_mass.shape, dtype=torch.int32).unsqueeze(0)
        self._atom_mask = self.atom_mask.numpy()
        # (A, 3)
        crds = mol2_obj.crds
        self.full_crds = torch.tensor(crds, dtype=torch.float32)
        # (B, 2)
        bond_index = mol2_obj.bond_indexes
        self.bonds = torch.tensor(bond_index).unsqueeze(0).to(torch.int32)
        # residue_atoms = mol2_obj.residue_atoms
        if index is None:
            raise ValueError('The rigid atom indexes should be set.')

        if residues is None:
            if isinstance(index, Tensor):
                self.rigid_index = index
            else:
                self.rigid_index = torch.tensor(index, dtype=torch.int32)
            self.tri_mask = self.rigid_index.reshape(-1)
            self.delete_mask = torch.from_numpy(np.setdiff1d(np.arange(self.full_crds.shape[-2]), self.tri_mask.numpy()))
            # (G, 3)
            group_rigid_index = self.rigid_index.reshape((self.rigid_index.shape[0], 3))
            # (G, 3)
            self.origins = (self.full_crds[group_rigid_index[:, 0]] + self.full_crds[group_rigid_index[:, 1]]) / 2
            group_atoms = int(self.full_crds.shape[-2] / (self.tri_mask.shape[0] / 3))
            # (G, g-3, 3)
            group_crds = self.full_crds[self.delete_mask].reshape((-1, group_atoms-3, 3))
            self.group_vec = group_crds - self.origins.unsqueeze(1)
        else:
            #TODO
            raise ValueError("To be implemented.")
        
        self.num_rigids = self.rigid_index.shape[-2]
        self.total_atoms = crds.shape[-2]
        self.num_atoms = self.total_atoms
        # (G, 3, 3)
        rcrd = self.full_crds[self.rigid_index]
        # (1, 4G, 3)
        self.Rt = self.rigid_from_points(rcrd).reshape((1, self.num_rigids*4, 3))
        # (1, G, 6)
        self.quat = Parameter(torch.zeros((1, self.num_rigids, 6), dtype=torch.float32), requires_grad=True)
        # (1, G, 3)
        self.standard_shape = (1, self.num_rigids, 3)
        self.build_system()
        self.build_space(self.full_crds, pbc_box=None)
    
    def hamiltonian_product(self, q, v):
        """Perform hamiltonian product.
        Args:
            q (Tensor): The transform quaternion.
            v (Tensor): The vector to be transformed.
        Returns:
            qvq^{-1}
        """
        # Get inverse quaternion
        iq = self.quaternion_inverse(q)
        # Perform hamiltonian product: q * v * q^{-1}
        op1 = self.quaternion_multiply(q, v)
        res = self.quaternion_multiply(op1, iq)
        return res
    
    def quaternion_multiply(self, q1, q2):
        """Calculate quaternion multiplication.
        Args:
            q1 (Tensor): First quaternion
            q2 (Tensor): Second quaternion
        Returns:
            Tensor: Product quaternion
        """
        w1, x1, y1, z1 = torch.unbind(q1, dim=-1)
        w2, x2, y2, z2 = torch.unbind(q2, dim=-1)
        
        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
        z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
        
        return torch.stack([w, x, y, z], dim=-1)
    
    def quaternion_inverse(self, q):
        """Calculate quaternion inverse.
        Args:
            q (Tensor): Input quaternion
        Returns:
            Tensor: Inverse quaternion
        """
        # Conjugate
        w, x, y, z = torch.unbind(q, dim=-1)
        conjugate = torch.stack([w, -x, -y, -z], dim=-1)
        
        # Normalize
        norm_squared = torch.sum(q * q, dim=-1, keepdim=True)
        inverse = conjugate / norm_squared
        
        return inverse
    
    def full_atoms(self):
        rigid_crds = self.get_rotate_rigid()
        full_crd = self.coordinate[0]
        # (G, 3, 3)
        rigid_xyz = torch.squeeze(self.get_xyz(rigid_crds))
        full_crd[self.tri_mask] = rigid_xyz.reshape(-1, 3)
        group_xyz = rigid_xyz.reshape((self.group_vec.shape[0], -1, 3))
        # (G, 3)
        new_origins = (group_xyz[:, 0] + group_xyz[:, 1]) / 2
        last_xyz = (self.calc_full_rotate() + new_origins.unsqueeze(0).unsqueeze(2)).reshape((-1, 3))
        full_crd[self.delete_mask] = last_xyz
        return full_crd.unsqueeze(0)
    
    def get_coordinate(self):
        return self.full_atoms()
    
    def rigid_from_points(self, crd):
        """ Transform the coordinates formulation. """
        # (N, 3)
        v1 = crd[:, 2] - crd[:, 1]
        v2 = crd[:, 0] - crd[:, 1]
        e1 = v1 / torch.norm(v1, dim=-1, keepdim=True)
        u2 = v2 - e1 * self.einsum_0(e1, v2).unsqueeze(-1)
        e2 = u2 / torch.norm(u2, dim=-1, keepdim=True)
        e3 = torch.cross(e1, e2, dim=-1)
        # (N, 3, 3)
        R = self.concat(e1.unsqueeze(1), e2.unsqueeze(1), e3.unsqueeze(1))
        t = crd[:, 1].unsqueeze(1)
        # (N, 4, 3)
        new_crd = self.concat(R, t)
        return new_crd

    def get_xyz(self, crd):
        """ Transform the (R,T) into cartesian coordinate system. """
        batches = crd.shape[0]
        # (B, N, 4, 3)
        rigid_crd = crd.reshape((batches, -1, 4, 3))
        # (B, N, 3)
        Ca = rigid_crd[:, :, -1]
        CaC = rigid_crd[:, :, 0] * self.L1
        C = Ca + CaC
        CaC_N = rigid_crd[:, :, 1] * self.Hight
        # N = Ca + NC
        N = C - rigid_crd[:, :, 0] * self.Floor + CaC_N
        # (B, N, 3, 3)
        xyz = torch.stack((N, Ca, C), dim=1).transpose(1, 2)
        # (B, 3N, 3)
        return xyz.reshape((batches, -1, 3))

    def get_rotate_rigid(self):
        """ Apply hamiltonian product. """
        # (B, N, 4, 3)
        triangle = self.Rt.reshape((self.coordinate.shape[0], -1, 4, 3))
        # (B, N, 4)
        e1 = F.pad(triangle[:, :, 0], (0, 0, 1, 0), mode='constant', value=0)
        e2 = F.pad(triangle[:, :, 1], (0, 0, 1, 0), mode='constant', value=0)
        e3 = F.pad(triangle[:, :, 2], (0, 0, 1, 0), mode='constant', value=0)
        # (B, N, 3)
        t = triangle[:, :, 3]
        # (B, N, 4)
        R = F.pad(self.quat[:, :, :3], (1, 0), mode='constant', value=1)
        R = R / torch.norm(R, dim=-1, keepdim=True)
        # (B, N, 3)
        dt = self.quat[:, :, 3:]
        t = t + dt
        # (B, N, 3)
        new_e1 = self.hamiltonian_product(R, e1)[:, :, 1:].reshape(self.standard_shape)
        new_e2 = self.hamiltonian_product(R, e2)[:, :, 1:].reshape(self.standard_shape)
        new_e3 = self.hamiltonian_product(R, e3)[:, :, 1:].reshape(self.standard_shape)
        # (B, N, 4, 3)
        new_triangle = torch.stack((new_e1, new_e2, new_e3, t), dim=1).transpose(1, 2).reshape(self.Rt.shape)
        return new_triangle
    
    def hamiltonian_product_2(self, R, crds):
        """Apply hamiltonian product to group of coordinates.
        
        Args:
            R (Tensor): Rotation quaternion
            crds (Tensor): Coordinates to transform
            
        Returns:
            Tensor: Transformed coordinates
        """
        # Reshape inputs for batch processing
        batch_size = R.shape[0]
        num_rigids = R.shape[1]
        num_points = crds.shape[2]
        
        # Reshape quaternions to match coordinates
        R_expanded = R.unsqueeze(2).expand(-1, -1, num_points, -1)
        R_flat = R_expanded.reshape(batch_size, -1, 4)
        
        # Reshape coordinates for processing
        crds_flat = crds.reshape(batch_size, -1, 4)
        
        # Apply hamiltonian product
        result = self.hamiltonian_product(R_flat, crds_flat)
        
        # Reshape back to original structure
        result = result.reshape(batch_size, num_rigids, num_points, 4)
        
        return result
    
    def calc_full_rotate(self):
        # (B, G, g-3, 4)
        crds = F.pad(self.group_vec, (1, 0), mode='constant', value=0).unsqueeze(0)
        # (B, G, 4)
        R = F.pad(self.quat[:, :, :3], (1, 0), mode='constant', value=1)
        R = R / torch.norm(R, dim=-1, keepdim=True)
        # (B, G, g-3, 3)
        new_crds = self.hamiltonian_product_2(R, crds).transpose(0, 1).reshape((self.coordinate.shape[0], self.quat.shape[1],
                                                                crds.shape[-2], crds.shape[-1]))[..., 1:]
        return new_crds
    
    def forward(self):
        pbc_box = None
        if self.pbc_box is not None:
            pbc_box = self.identity(self.pbc_box)
        return self.full_atoms(), pbc_box
